from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse, HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from kokoro import KModel, KPipeline
from pathlib import Path
import numpy as np
import soundfile as sf
import torch
import io
import json
import sqlite3
import datetime
import os
from typing import Optional, List, Dict

# 初始化FastAPI应用
app = FastAPI(title="Kokoro TTS 流式语音服务")

# 数据库路径
DB_PATH = "tts_feedback.db"

# 初始化数据库
def init_db():
    if not os.path.exists(DB_PATH):
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        
        # 创建点赞/点踩表
        cursor.execute('''
        CREATE TABLE feedback (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            text_id TEXT,
            rating INTEGER,
            content TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
        ''')
        
        # 创建访问数据表
        cursor.execute('''
        CREATE TABLE access_log (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            ip TEXT,
            text TEXT,
            voice TEXT,
            speed REAL,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
        ''')
        
        conn.commit()
        conn.close()
        print("数据库初始化成功")

# 初始化数据库
init_db()

# 点赞/点踩请求模型
class FeedbackRequest(BaseModel):
    text_id: str
    rating: int  # 1表示点赞，-1表示点踩
    content: str = ""  # 反馈内容

# 访问数据请求模型
class AccessLogRequest(BaseModel):
    ip: str
    text: str
    voice: str
    speed: float

# 允许跨域请求（方便前端调用）
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 生产环境需指定具体域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 模型配置
REPO_ID = 'hexgrad/Kokoro-82M-v1.1-zh'
SAMPLE_RATE = 24000  # 采样率固定为24000Hz
N_ZEROS = 5000  # 段落间静音长度（约0.2秒）
device = 'cuda' if torch.cuda.is_available() else 'cpu'

en_pipeline = KPipeline(lang_code='a', repo_id=REPO_ID, model=False)
def en_callable(text):
    if text == 'Kokoro':
        return 'kˈOkəɹO'
    elif text == 'Sol':
        return 'sˈOl'
    return next(en_pipeline(text)).phonemes

# 加载模型（全局单例，避免重复加载）
class TTSModel:
    def __init__(self):
        self.model = None
        self.zh_pipeline = None
        self.voices = {
            "zf_001": "./voices/zf_001.pt",  # 女性音色
            "zm_010": "./voices/zm_010.pt"   # 男性音色（若有）
        }
        
        self.loaded = False

    def load(self, model_path="./kokoro-v1_1-zh.pth", config_path="./config.json"):
        """加载模型（首次调用时自动加载）"""
    #     def en_callable(text):
    #         en_pipeline = KPipeline(lang_code='a', repo_id=REPO_ID, model=self.model)
    #         if text == 'Kokoro':
    #             return 'kˈOkəɹO'
    #         elif text == 'Sol':
    #             return 'sˈOl'
    #         return {"phonemes":next(en_pipeline(text)).phonemes,
    #     "voice": "./voices/af_heart.pt" ,  # 默认英文音色
    # }
        if not self.loaded:
            try:
                self.model = KModel(
                    config=config_path,
                    repo_id=REPO_ID,
                    model=model_path
                ).to(device).eval()
                self.zh_pipeline = KPipeline(lang_code='z', repo_id=REPO_ID, model=self.model,en_callable=en_callable)
                self.loaded = True
                print(f"模型加载成功，使用设备：{device}")
            except Exception as e:
                print(f"模型加载失败：{e}")
                raise HTTPException(status_code=500, detail=f"模型初始化失败：{str(e)}")

# 实例化模型（全局唯一）
tts_model = TTSModel()

# 请求参数模型
class TTSRequest(BaseModel):
    text: str  # 需要合成的文本
    voice: Optional[str] = "zf_001"  # 音色选择
    speed: Optional[float] = 1.0  # 语速（0.5-2.0）

# 语速调整函数（根据文本长度动态调整）
def speed_callable(len_ps, base_speed=1.0):
    """
    动态调整语速：文本越长，语速越慢（避免合成过于急促）
    :param len_ps: 文本长度
    :param base_speed: 基础语速（用户指定）
    """
    speed = 0.8
    if len_ps <= 83:
        speed = 1.0
    elif len_ps < 183:
        speed = 1.0 - (len_ps - 83) / 500
    return speed * base_speed * 1.1  # 保持与原代码一致的系数

# 流式生成音频的生成器函数
def tts_stream_generator(text: str, voice: str, speed: float):
    """
    流式生成音频片段，逐句返回WAV格式数据
    """
    # 确保模型已加载
    tts_model.load()
    
    # 检查音色是否存在
    if voice not in tts_model.voices:
        raise HTTPException(status_code=400, detail=f"不支持的音色：{voice}，可选音色：{list(tts_model.voices.keys())}")
    voice_path = tts_model.voices[voice]
    
    # 简单分句（按中文标点分割）
    # 优化文本拆分逻辑，使用正则表达式更精确地匹配标点符号及后续空白
    import re
    sentences = [s.strip() for s in re.split(r'[，。,;；？！?！]\s*', text) if s.strip()]
    print(len(sentences))   
    if not sentences:
        raise HTTPException(status_code=400, detail="输入文本为空或无法分割")
    
    # 逐句生成音频
    first_sentence = True

    i=0
    wav_data=[]
    #便利sentences,如果长度小于50，则跟下一个拼接到一起，直到每个元素的长度都大于50，允许最后一个元素小于50
    #如果sentences只有一句则啥也不做
    #处理一下各种情况可能存在的边界问题
    
    if len(sentences)>1:
        while i<len(sentences):
            if len(sentences[i])<50:
                #如果i+1不存在，直接break
                if i+1>=len(sentences):
                    break
                #拼接加一个逗号
                sentences[i]+='，'
                sentences[i]+=sentences[i+1]
                sentences.pop(i+1)
            else:
                i+=1
    i=0
    for sentence in sentences:
        try:
            print(sentence)
            # 生成当前句的音频
            generator = tts_model.zh_pipeline(
                sentence,
                voice=voice_path,
                speed=lambda len_ps: speed_callable(len_ps, base_speed=speed)

            )
            result = next(generator)
            wav = result.audio  # 音频数据（numpy数组）

            # 非首句前添加静音
            if not first_sentence and N_ZEROS > 0:
                silence = np.zeros(N_ZEROS, dtype=np.float32)
                wav = np.concatenate([silence, wav])
            
            # first_sentence = False
            
            buffer = io.BytesIO()
            if first_sentence:
                first_sentence=False
                sf.write(buffer, wav, SAMPLE_RATE, format="WAV", subtype="PCM_16")
            else:
                sf.write(buffer, wav, SAMPLE_RATE, format="WAV", subtype="PCM_16")
            buffer.seek(0)
            #将buffer保存一份为本地文件，并yield出去
            # wav_data.append(buffer.read())
            
            # buffer.seek(0)
            yield buffer.read()  # 流式返回当前片段
            
        except Exception as e:
            print(f"生成音频失败（句子：{sentence}）：{e}")
            raise HTTPException(status_code=500, detail=f"生成音频失败：{str(e)}")
    #将wav_data生成到一个wav里面
    # with sf.SoundFile('stream_output.wav', 'w', SAMPLE_RATE, channels=1) as f:  # 1声道
    #     import numpy as np
    #     for block in wav_data:
    #         # Convert bytes back to numpy array with correct dtype (int16 for PCM_16)
    #         audio_data = np.frombuffer(block, dtype=np.int16)
    #         f.write(audio_data)
# TTS流式接口
@app.post("/api/tts/stream", summary="流式语音合成接口")
async def tts_stream(request: TTSRequest):
    """
    流式返回语音合成结果（WAV格式）
    - 支持逐句生成并返回音频，前端可边接收边播放
    - 音频格式：PCM_16，采样率24000Hz，单声道
    """
    try:
        # 返回流式响应，MIME类型为音频流
        return StreamingResponse(
            tts_stream_generator(request.text.replace("\n",","), request.voice, request.speed),
            media_type="audio/wav",
            headers={
                "Content-Disposition": 'attachment; filename="tts_stream.wav"',
                "Transfer-Encoding": "chunked",  # 分块传输
                "Content-Type": "application/octet-stream",
                "Connection": "keep-alive"
            }
        )
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"服务错误：{str(e)}")

# 点赞/点踩接口
@app.post("/api/feedback", summary="点赞/点踩反馈接口")
async def submit_feedback(feedback: FeedbackRequest):
    """
    提交点赞/点踩反馈
    - text_id: 文本ID
    - rating: 评分（1表示点赞，-1表示点踩）
    - content: 反馈内容（可选）
    """
    try:
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        cursor.execute(
            "INSERT INTO feedback (text_id, rating, content, created_at) VALUES (?, ?, ?, ?)",
            (feedback.text_id, feedback.rating, feedback.content, datetime.datetime.now().isoformat())
        )
        conn.commit()
        conn.close()
        return {"status": "success", "message": "反馈提交成功"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"提交反馈失败：{str(e)}")

# 访问数据收集接口
@app.post("/api/access-log", summary="访问数据收集接口")
async def submit_access_log(log: AccessLogRequest):
    """
    提交访问数据
    - ip: IP地址
    - text: 输入文本
    - voice: 选择的音色
    - speed: 语速
    """
    try:
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        cursor.execute(
            "INSERT INTO access_log (ip, text, voice, speed, created_at) VALUES (?, ?, ?, ?, ?)",
            (log.ip, log.text, log.voice, log.speed, datetime.datetime.now().isoformat())
        )
        conn.commit()
        conn.close()
        return {"status": "success", "message": "访问数据记录成功"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"记录访问数据失败：{str(e)}")

# 获取反馈统计接口
@app.get("/api/feedback/stats", summary="获取反馈统计数据")
async def get_feedback_stats():
    """
    获取点赞/点踩统计数据
    """
    try:
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        
        # 获取点赞总数
        cursor.execute("SELECT COUNT(*) FROM feedback WHERE rating = 1")
        likes = cursor.fetchone()[0]
        
        # 获取点踩总数
        cursor.execute("SELECT COUNT(*) FROM feedback WHERE rating = -1")
        dislikes = cursor.fetchone()[0]
        
        conn.close()
        
        return {
            "status": "success",
            "data": {
                "likes": likes,
                "dislikes": dislikes,
                "total": likes + dislikes
            }
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取统计数据失败：{str(e)}")

# 获取最近访问记录接口
@app.get("/api/access-log/recent", summary="获取最近访问记录")
async def get_recent_access_logs(limit: int = 10):
    """
    获取最近访问记录
    - limit: 返回记录数量，默认为10
    """
    try:
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        
        cursor.execute(
            "SELECT ip, text, voice, speed, created_at FROM access_log ORDER BY created_at DESC LIMIT ?",
            (limit,)
        )
        logs = cursor.fetchall()
        
        conn.close()
        
        # 转换为字典列表
        result = []
        for log in logs:
            result.append({
                "ip": log[0],
                "text": log[1],
                "voice": log[2],
                "speed": log[3],
                "created_at": log[4]
            })
        
        return {
            "status": "success",
            "data": result
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取访问记录失败：{str(e)}")

# 首页（提供简单的测试界面）
@app.get("/", response_class=HTMLResponse)
async def index():
    return """
    <html>
        <head>
            <title>Kokoro TTS 流式服务</title>
            <style>
                body { font-family: Arial, sans-serif; max-width: 800px; margin: 0 auto; padding: 20px; }
                .container { text-align: center; }
                textarea { width: 100%; height: 150px; margin: 20px 0; padding: 10px; }
                .controls { margin: 10px 0; }
                button { padding: 10px 20px; background: #165DFF; color: white; border: none; border-radius: 5px; cursor: pointer; }
                #status { color: #666; margin: 10px 0; }
            </style>
        </head>
        <body>
            <div class="container">
                <h1>Kokoro TTS 流式语音合成</h1>
                <textarea id="textInput">Kokoro 是一系列体积虽小但功能强大的 TTS 模型。该模型是经过短期训练的结果。</textarea>
                <div class="controls">
                    <label>音色：</label>
                    <select id="voiceSelect">
                        <option value="zf_001">女性音色 (zf_001)</option>
                        <option value="zm_010">男性音色 (zm_010)</option>
                    </select>
                    <label style="margin-left: 20px;">语速：</label>
                    <input type="range" id="speedRange" min="0.5" max="2.0" step="0.1" value="1.0">
                    <span id="speedValue">1.0</span>
                </div>
                <button onclick="startTTS()">开始合成并播放</button>
                <div id="status">就绪</div>
                <audio id="audioPlayer" controls style="width: 100%; margin-top: 20px;"></audio>
            </div>
            <script>
                // 显示当前语速值
                document.getElementById('speedRange').addEventListener('input', function() {
                    document.getElementById('speedValue').textContent = this.value;
                });
                
                // 开始TTS合成并播放
                async function startTTS() {
                    const text = document.getElementById('textInput').value;
                    const voice = document.getElementById('voiceSelect').value;
                    const speed = parseFloat(document.getElementById('speedRange').value);
                    const status = document.getElementById('status');
                    const audioPlayer = document.getElementById('audioPlayer');
                    
                    if (!text.trim()) {
                        alert('请输入文本');
                        return;
                    }
                    
                    status.textContent = '正在合成语音...';
                    
                    try {
                        // 调用流式接口
                        const response = await fetch('/api/tts/stream', {
                            method: 'POST',
                            headers: { 'Content-Type': 'application/json' },
                            body: JSON.stringify({ text, voice, speed })
                        });
                        
                        if (!response.ok) {
                            throw new Error(`HTTP错误：${response.status}`);
                        }
                        
                        // 使用MediaSource实现流式播放
                        const mediaSource = new MediaSource();
                        audioPlayer.src = URL.createObjectURL(mediaSource);

                        mediaSource.addEventListener('sourceopen', async () => {
                            const sourceBuffer = mediaSource.addSourceBuffer('audio/mpeg; codecs="1"');
                            const reader = response.body.getReader();
                            status.textContent = '正在接收并播放音频...';

                            while (true) {
                                const { done, value } = await reader.read();
                                if (done) {
                                    if (!sourceBuffer.updating) {
                                    mediaSource.endOfStream();
                                } else {
                                    sourceBuffer.addEventListener('updateend', function onUpdateEnd() {
                                        sourceBuffer.removeEventListener('updateend', onUpdateEnd);
                                        mediaSource.endOfStream();
                                    });
                                }
                                    status.textContent = '合成播放完成';
                                    break;
                                }

                                // 等待sourceBuffer可以接受新数据
                                while (sourceBuffer.updating) {
                                    await new Promise(resolve => sourceBuffer.addEventListener('updateend', resolve, { once: true }));
                                }

                                // 将接收到的数据添加到sourceBuffer
                                sourceBuffer.appendBuffer(value);
                            }
                        });

                        // 当媒体可以播放时自动播放
                        audioPlayer.addEventListener('canplay', () => {
                            if (status.textContent.includes('正在接收')) {
                                audioPlayer.play().catch(e => console.error('播放失败:', e));
                            }
                        });

                    status.textContent = '合成完成';
                        
                    } catch (error) {
                        status.textContent = `错误：${error.message}`;
                        console.error('TTS失败：', error);
                    }
                }
            </script>
        </body>
    </html>
    """

# 获取最近的反馈记录
@app.get("/api/feedback/recent")
async def get_recent_feedback(limit: int = 100):
    """获取最近的反馈记录"""
    try:
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        
        cursor.execute(
            "SELECT text_id, rating, content, created_at FROM feedback ORDER BY created_at DESC LIMIT ?",
            (limit,)
        )
        feedbacks = cursor.fetchall()
        
        conn.close()
        
        # 转换为字典列表
        result = []
        for feedback in feedbacks:
            result.append({
                "text_id": feedback[0],
                "rating": feedback[1],
                "content": feedback[2],
                "created_at": feedback[3]
            })
        
        return {
            "status": "success",
            "data": result
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取反馈记录失败：{str(e)}")

# 启动服务
if __name__ == "__main__":
    import uvicorn
    # 启动服务（host=0.0.0.0 允许外部访问，port=8000 端口）
    uvicorn.run(app, host="0.0.0.0", port=8000, workers=1)  # TTS模型不支持多进程，workers=1